
#include "reflection_pad1d_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
// 常量定义
const uint32_t BYTE_BLOCK = 32;                   // 32字节对齐
const int32_t X_INPUT_INDEX = 0;                  // 输入张量索引
const int32_t PAD_INPUT_INDEX = 1;                // 填充参数索引
const int32_t OUTPUT_INDEX = 0;                   // 输出张量索引

const int32_t FLOAT_BYTES = 4;                    // FP32字节数
const int32_t FLOAT16_BYTES = 2;                  // FP16/BF16字节数

// 维度索引映射
const uint32_t DIM_1D_W = 0;                      // 1D: [W]
const uint32_t DIM_2D_C = 0;                      // 2D: [C, W]
const uint32_t DIM_2D_W = 1;
const uint32_t DIM_3D_N = 0;                      // 3D: [N, C, W]
const uint32_t DIM_3D_C = 1;
const uint32_t DIM_3D_W = 2;

// 填充参数索引
const uint32_t PAD_LEFT_INDEX = 0;                // 左填充
const uint32_t PAD_RIGHT_INDEX = 1;               // 右填充
const uint32_t PAD_LEN = 2;                       // 填充参数长度

// 硬件相关常量
const uint32_t RESERVED_UB = 32 * 1024;           // 预留UB内存(32KB)
const uint32_t ALIGN_256_BYTES = 256;             // UB对齐粒度
const uint32_t MAX_CORE_NUM = 64;                  // 最大核心数
const uint32_t UB_PARTITION_NUM = 5;               // UB分区数
const uint32_t MIN_UB_ELEMENTS = 32;               // 最小UB元素数

// TilingKey定义
const uint32_t TILING_KEY_FP32 = 0;
const uint32_t TILING_KEY_FP16 = 1;
const uint32_t TILING_KEY_BF16 = 2;

// 数据类型字节映射表
static std::map<ge::DataType, int32_t> DATATYPE_BYTES_MAP = {
    {ge::DT_FLOAT, FLOAT_BYTES},
    {ge::DT_FLOAT16, FLOAT16_BYTES},
    {ge::DT_BF16, FLOAT16_BYTES}
};

// 数据类型TilingKey映射表
static std::map<ge::DataType, uint32_t> DTYPE_TILING_KEY_MAP = {
    {ge::DT_FLOAT, TILING_KEY_FP32},
    {ge::DT_FLOAT16, TILING_KEY_FP16},
    {ge::DT_BF16, TILING_KEY_BF16}
};

// 向上对齐函数
template <typename T1, typename T2>
static inline T1 CeilAlign(T1 value, T2 align) {
    if (align == 0) return value;
    return (value + align - 1) / align * align;
}

// 向下对齐函数
template <typename T1, typename T2>
static inline T1 FloorAlign(T1 value, T2 align) {
    if (align == 0) return value;
    return value / align * align;
}

// 主Tiling函数
static ge::graphStatus TilingFunc(gert::TilingContext* context) {
    if (context == nullptr) {
        return ge::GRAPH_FAILED;
    }

    // 获取编译信息
    uint32_t coreNum = 0;
    uint64_t ubSizePlatForm = 0;
    // uint32_t sysWorkspaceSize = 0;
    platform_ascendc::PlatformAscendC ascendcPlatform(context->GetPlatformInfo());

    // 获取AIV核心数
    coreNum = ascendcPlatform.GetCoreNum();
    if (coreNum == 0) {
        coreNum = 1;
    } else if (coreNum > MAX_CORE_NUM) {
        coreNum = MAX_CORE_NUM;
    }

    // 获取UB内存大小
    uint64_t ubByteSize = 0;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubByteSize);
    if (ubByteSize == 0) {
        return ge::GRAPH_FAILED;
    }
    ubSizePlatForm = ubByteSize;

    // 获取系统工作空间
    // sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();

    // 解析输入张量
    const auto inputTensor = context->GetInputTensor(X_INPUT_INDEX);
    if (inputTensor == nullptr) {
        return ge::GRAPH_FAILED;
    }
    ge::DataType inputDtype = inputTensor->GetDataType();


    // 校验数据类型
    if (DATATYPE_BYTES_MAP.find(inputDtype) == DATATYPE_BYTES_MAP.end()) {
        return ge::GRAPH_FAILED;
    }
    int32_t dtypeBytes = DATATYPE_BYTES_MAP[inputDtype];

    // 解析输入形状
    const auto inputShape = context->GetInputShape(X_INPUT_INDEX)->GetOriginShape();
    uint32_t inputDim = static_cast<uint32_t>(inputShape.GetDimNum());
    if (inputDim < 1 || inputDim > 3) {
        return ge::GRAPH_FAILED;
    }

    // 映射N/C/W维度
    uint32_t nSize = 1, cSize = 1, wSize = 1;
    switch (inputDim) {
        case 1:
            wSize = static_cast<uint32_t>(inputShape[DIM_1D_W]);
            break;
        case 2:
            cSize = static_cast<uint32_t>(inputShape[DIM_2D_C]);
            wSize = static_cast<uint32_t>(inputShape[DIM_2D_W]);
            break;
        case 3:
            nSize = static_cast<uint32_t>(inputShape[DIM_3D_N]);
            cSize = static_cast<uint32_t>(inputShape[DIM_3D_C]);
            wSize = static_cast<uint32_t>(inputShape[DIM_3D_W]);
            break;
        default:
            return ge::GRAPH_FAILED;
    }

    // 解析填充参数
    const auto paddingTensor = context->GetInputTensor(PAD_INPUT_INDEX);
    if (paddingTensor == nullptr) {
        return ge::GRAPH_FAILED;
    }
    uint32_t padLeft = 0, padRight = 0;

    if (paddingTensor->GetDataType() == ge::DT_INT32) {
        const int32_t* padData = paddingTensor->GetData<int32_t>();
        if (padData == nullptr) {
            return ge::GRAPH_FAILED;
        }

        uint64_t padLen = paddingTensor->GetShapeSize();
        if (padLen != PAD_LEN) {
            return ge::GRAPH_FAILED;
        }

        int64_t left = static_cast<int64_t>(padData[PAD_LEFT_INDEX]);
        int64_t right = static_cast<int64_t>(padData[PAD_RIGHT_INDEX]);
        if (left < 0 || right < 0) {
            return ge::GRAPH_FAILED;
        }

        padLeft = static_cast<uint32_t>(left);
        padRight = static_cast<uint32_t>(right);
    } else if (paddingTensor->GetDataType() == ge::DT_INT64) {
        const int64_t* padData = paddingTensor->GetData<int64_t>();
        if (padData == nullptr) {
            return ge::GRAPH_FAILED;
        }

        uint64_t padLen = paddingTensor->GetShapeSize();
        if (padLen != PAD_LEN) {
            return ge::GRAPH_FAILED;
        }

        int64_t left = static_cast<int64_t>(padData[PAD_LEFT_INDEX]);
        int64_t right = static_cast<int64_t>(padData[PAD_RIGHT_INDEX]);
        if (left < 0 || right < 0) {
            return ge::GRAPH_FAILED;
        }

        padLeft = static_cast<uint32_t>(left);
        padRight = static_cast<uint32_t>(right);
    } else {
        return ge::GRAPH_FAILED;
    }

    // 计算输出形状
    uint32_t outWSize = wSize + padLeft + padRight;
    uint32_t elePer32B = BYTE_BLOCK / dtypeBytes;
    uint32_t alignWSize = CeilAlign(wSize, elePer32B);
    uint32_t alignOutWSize = CeilAlign(outWSize, elePer32B);

    // 分核策略计算
    uint64_t totalNc = static_cast<uint64_t>(nSize) * cSize;
    uint32_t blockNum = 0, ncPerCore = 0, tailNC = 0;

    if (totalNc == 0 || coreNum == 0) {
        return ge::GRAPH_FAILED;
    }

    if (totalNc <= coreNum) {
        blockNum = static_cast<uint32_t>(totalNc);
        ncPerCore = 1;
        tailNC = 0;
    } else {
        ncPerCore = static_cast<uint32_t>(totalNc / coreNum);
        tailNC = static_cast<uint32_t>(totalNc % coreNum);
        blockNum = coreNum;
    }

    // UB内存分配
    uint32_t tilingDataSize = CeilAlign(sizeof(ReflectionPad1dTilingData), BYTE_BLOCK);
    if (ubSizePlatForm < RESERVED_UB + tilingDataSize) {
        return ge::GRAPH_FAILED;
    }
    
    uint32_t availableUb = 0;
    uint64_t availableUb64 = FloorAlign(ubSizePlatForm - RESERVED_UB - tilingDataSize, BYTE_BLOCK);
    if (availableUb64 > UINT32_MAX) {
        availableUb = UINT32_MAX;
    } 
    else {
        availableUb = static_cast<uint32_t>(availableUb64);
    }
    
    uint32_t ubPerPartition = availableUb / UB_PARTITION_NUM;
    uint32_t ubElementNum = ubPerPartition / dtypeBytes;
    ubElementNum = FloorAlign(ubElementNum, ALIGN_256_BYTES / dtypeBytes);

    // 确保UB元素数不为0
    if (ubElementNum < MIN_UB_ELEMENTS) {
        ubElementNum = MIN_UB_ELEMENTS;
    }
    // 工作空间计算
    uint64_t workspacePerCore = 0;
    // bool isLargePad = (padLeft + padRight > ubElementNum / 2) ||
    //                  (static_cast<uint64_t>(wSize) + padLeft + padRight > ubElementNum);

    // if ((padLeft > 0 || padRight > 0) && isLargePad) {
    //     workspacePerCore = static_cast<uint64_t>(outWSize) * ncPerCore * dtypeBytes;
    //     if (tailNC > 0) {
    //         workspacePerCore = std::max(workspacePerCore,
    //             static_cast<uint64_t>(outWSize) * (ncPerCore + 1) * dtypeBytes);
    //     }
    //     workspacePerCore = CeilAlign(workspacePerCore, BYTE_BLOCK);
    // }

    // 填充TilingData
    ReflectionPad1dTilingData tilingData;
    tilingData.set_wSize(wSize);
    tilingData.set_alignWSize(alignWSize);
    tilingData.set_padLeft(padLeft);
    tilingData.set_padRight(padRight);
    tilingData.set_blockNum(blockNum);
    tilingData.set_ncPerCore(ncPerCore);
    tilingData.set_tailNC(tailNC);
    tilingData.set_ubElementNum(ubElementNum);
    tilingData.set_tilingKey(DTYPE_TILING_KEY_MAP[inputDtype]);
    tilingData.set_workspacePerCore(workspacePerCore);

    // 设置Kernel配置
    context->SetBlockDim(blockNum);
    context->SetTilingKey(tilingData.get_tilingKey());

    // 计算总工作空间
    size_t* workspaceSizes = context->GetWorkspaceSizes(1);
    // if (workspaceSizes == nullptr) {
         
    //     return ge::GRAPH_FAILED;
    // }
    // workspaceSizes[0] = workspacePerCore * blockNum + sysWorkspaceSize;
    workspaceSizes[0] = 0;

    // 序列化TilingData
    auto rawTilingData = context->GetRawTilingData();
    tilingData.SaveToBuffer(rawTilingData->GetData(), rawTilingData->GetCapacity());
    rawTilingData->SetDataSize(tilingData.GetDataSize());
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge {
const size_t IDX_IN_PADDING = 1;
const uint32_t PAD_W_FRONT_INDEX = 0;             // W前填充（padLeft）
const uint32_t PAD_W_BACK_INDEX = 1;              // W后填充（padRight）

static ge::graphStatus InferShape(gert::InferShapeContext* context) {
    if (context == nullptr) {
        return ge::GRAPH_FAILED;
    }

    // 获取输入形状
    const gert::Shape* inputShape = context->GetInputShape(0);  // 输入x的形状
    if (inputShape == nullptr) {
        return ge::GRAPH_FAILED;
    }
    uint32_t inputDim = static_cast<uint32_t>(inputShape->GetDimNum());
    if (inputDim < 1 || inputDim > 3) {  // 限制输入维度为1D~3D
        return ge::GRAPH_FAILED;
    }

    // 获取paddings参数
    const auto paddingTensor = context->GetInputTensor(1);  // 输入paddings的张量
    if (paddingTensor == nullptr) {
        return ge::GRAPH_FAILED;
    }
    const gert::Shape* padShape = context->GetInputShape(1);
    if (padShape == nullptr || padShape->GetDimNum() != 1 || padShape->GetDim(0) != 2) {
        return ge::GRAPH_FAILED; 
    }

    ge::DataType padDtype = paddingTensor->GetDataType();
    uint32_t padLeft = 0, padRight = 0;
    if (padDtype == ge::DT_INT32) {
        const int32_t* padData = paddingTensor->GetData<int32_t>();
        if (padData == nullptr) return ge::GRAPH_FAILED;
        padLeft = static_cast<uint32_t>(padData[PAD_W_FRONT_INDEX]); 
        padRight = static_cast<uint32_t>(padData[PAD_W_BACK_INDEX]); 
    } 
    else if (padDtype == ge::DT_INT64) {
        const int64_t* padData = paddingTensor->GetData<int64_t>();
        if (padData == nullptr) return ge::GRAPH_FAILED;
        padLeft = static_cast<uint32_t>(padData[PAD_W_FRONT_INDEX]);
        padRight = static_cast<uint32_t>(padData[PAD_W_BACK_INDEX]);
    } 
    else {
        return ge::GRAPH_FAILED;  // 仅支持int32/int64类型paddings
    }

    if (padLeft > static_cast<uint32_t>(INT32_MAX) || padRight > static_cast<uint32_t>(INT32_MAX)) {
        return ge::GRAPH_FAILED;
    }

    // 计算输出形状
    gert::Shape* outputShape = context->GetOutputShape(0);
    if (outputShape == nullptr) {
        return ge::GRAPH_FAILED;
    }
    *outputShape = *inputShape; 
    uint32_t wDimIndex = static_cast<uint32_t>(inputShape->GetDimNum() - 1);  
    uint32_t inWSize = static_cast<uint32_t>(inputShape->GetDim(wDimIndex));
    
    if (static_cast<uint64_t>(inWSize) + padLeft + padRight > static_cast<uint64_t>(UINT32_MAX)) {
        return ge::GRAPH_FAILED;
    }
    uint32_t outWSize = inWSize + padLeft + padRight;
    outputShape->SetDim(wDimIndex, outWSize);  

    return ge::GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class ReflectionPad1d : public OpDef {
public:
    explicit ReflectionPad1d(const char* name) : OpDef(name) {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});  

        this->Input("paddings")
            .ParamType(REQUIRED)
            .ValueDepend(REQUIRED)
            .DataTypeList({ge::DT_INT64})
            .FormatList({ge::FORMAT_ND});

        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore()
            .AddConfig("ascend910")	
            .AddConfig("ascend310p")	
            .AddConfig("ascend310b")	
            .AddConfig("ascend910b");
    }
};

OP_ADD(ReflectionPad1d);
} // namespace ops